package com.atguigu.aiproject.controller.ai;
import com.atguigu.aiproject.agent.TweetGenerationAgent;
import com.atguigu.aiproject.config.SseConnectionManager;
import jakarta.servlet.http.HttpServletRequest;
import lombok.extern.slf4j.Slf4j;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.beans.factory.annotation.Qualifier;
import org.springframework.core.task.TaskExecutor;
import org.springframework.http.MediaType;
import org.springframework.web.bind.annotation.*;
import org.springframework.web.servlet.mvc.method.annotation.SseEmitter;
import reactor.core.publisher.Flux;
import java.io.IOException;
import java.util.List;
import java.util.concurrent.atomic.AtomicBoolean;

@RestController
@RequestMapping("/ai/tweet")
@Slf4j
public class TweetController {

    @Autowired
    private TweetGenerationAgent tweetGenerationAgent;//智能体

    @Autowired
    @Qualifier("sseTaskExecutor")
    private TaskExecutor taskExecutor;//线程池

    @Autowired
    private SseConnectionManager connectionManager;//sse连接管理

    /**
     * 流式生成推文和图片
     */
    @GetMapping(path = "/stream-images", produces = MediaType.TEXT_EVENT_STREAM_VALUE)
    public SseEmitter streamTweetWithImages(
            @RequestParam String topic,
            @RequestParam(defaultValue = "2") int imageCount,
            HttpServletRequest request) {

        // 创建带超时的Emitter
        SseEmitter emitter = new SseEmitter(30_000L);
        String sessionId = generateSessionId(request);

        // 注册连接管理
        connectionManager.addConnection(sessionId, emitter);

        // 设置事件处理器
        setupEmitterHandlers(emitter, sessionId);

        // 异步处理业务逻辑
        taskExecutor.execute(() -> processTweetGeneration(emitter, topic, imageCount, sessionId));

        return emitter;
    }

    private void setupEmitterHandlers(SseEmitter emitter, String sessionId) {
        emitter.onTimeout(() -> {
            log.info("SSE连接超时, session: {}", sessionId);
            connectionManager.removeConnection(sessionId);
        });

        emitter.onCompletion(() -> {
            log.info("SSE连接完成, session: {}", sessionId);
            connectionManager.removeConnection(sessionId);
        });

        emitter.onError(error -> {
            log.error("SSE连接错误, session: {}", sessionId, error);
            connectionManager.removeConnection(sessionId);
        });
    }

    /**
     * 处理推文生成和图片生成
     */
    private void processTweetGeneration(SseEmitter emitter, String topic, int imageCount, String sessionId) {
        try {
            AtomicBoolean isCompleted = new AtomicBoolean(false);

            // 流式生成推文
            tweetGenerationAgent.generateTweetStream(topic)
                    .doOnNext(content -> sendSseEvent(emitter, "tweet", content, isCompleted))
                    .doOnComplete(() -> handleTweetComplete(emitter, topic, imageCount, isCompleted))
                    .doOnError(error -> handleGenerationError(emitter, error, isCompleted))
                    .subscribe();

        } catch (Exception e) {
            handleGenerationError(emitter, e, new AtomicBoolean(false));
        }
    }

    // 处理推文完成，生成图片
    private void handleTweetComplete(SseEmitter emitter, String topic, int imageCount, AtomicBoolean isCompleted) {
        if (isCompleted.get()) return;

        try {
            sendSseEvent(emitter, "tweetComplete", "", isCompleted);

            if (imageCount <= 0) {
                emitter.complete();
                return;
            }

            // 生成图片提示词和图片
            String imagePrompt = tweetGenerationAgent.generateImagePrompt(topic);//根据主题生成图片提示词
            log.info("图片提示词: {}", imagePrompt);

            if (imageCount == 1) {
                String imageResult = tweetGenerationAgent.generateImage(imagePrompt);
                sendSseEvent(emitter, "imageResult", imageResult, isCompleted);
            } else {
                List<String> imageResults = tweetGenerationAgent.generateImages(imagePrompt, imageCount);
                sendSseEvent(emitter, "imageResults", imageResults, isCompleted);
            }

            emitter.complete();

        } catch (Exception e) {
            handleGenerationError(emitter, e, isCompleted);
        }
    }

    private void sendSseEvent(SseEmitter emitter, String eventName, Object data, AtomicBoolean isCompleted) {
        if (isCompleted.get()) return;

        try {
            emitter.send(SseEmitter.event().name(eventName).data(data));
        } catch (IOException e) {
            log.warn("发送SSE事件失败，连接可能已关闭", e);
            isCompleted.set(true);
        }
    }

    private void handleGenerationError(SseEmitter emitter, Throwable error, AtomicBoolean isCompleted) {
        if (isCompleted.getAndSet(true)) return;

        try {
            emitter.send(SseEmitter.event()
                    .name("error")
                    .data("生成失败: " + error.getMessage()));
            emitter.completeWithError(error);
        } catch (Exception e) {
            log.error("发送错误事件失败", e);
        }
    }

    private String generateSessionId(HttpServletRequest request) {
        return request.getSession().getId() + "-" + System.currentTimeMillis();
    }
}